import os
import torch
import torch.nn as nn
import numpy as np
import joblib
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score

matplotlib.use('Agg')
device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# ------- 定义模型结构 -------
class SlimModelCNN(nn.Module):
    def __init__(self, ch1, ch2):
        super(SlimModelCNN, self).__init__()
        self.conv1 = nn.Conv1d(15, ch1, kernel_size=3, padding=1)
        self.pool1 = nn.MaxPool1d(2)
        self.conv2 = nn.Conv1d(ch1, ch2, kernel_size=3, padding=1)
        self.pool2 = nn.MaxPool1d(2)
        self.conv3 = nn.Conv1d(ch2, 1, kernel_size=3, padding=1)
        self.conv4 = nn.Conv1d(32, 1, kernel_size=1)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = x.permute(0, 2, 1)
        x = self.conv1(x)
        x = self.relu(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.pool2(x)
        x = self.conv3(x)
        x = x.permute(0, 2, 1)
        x = self.conv4(x)
        x = x.squeeze(1)
        return x

def load_model(model, path):
    model.load_state_dict(torch.load(path))
    print(f"Model loaded from {path}")

def preprocess_data(data, save_dir, is_train=False):
    # 滑窗，归一化，与原版保持一致
    def create_sliding_window(data, window_size=128, step_size=64):
        X, y, indices = [], [], []
        for i in range(0, len(data) - window_size, step_size):
            window = data[i:i + window_size]
            features = window[:, 1:]  # 去掉第一列
            target = window[-1, 0]    # 最后一帧的目标值
            X.append(features)
            y.append(target)
            indices.append(i + window_size - 1)
        return np.array(X), np.array(y), np.array(indices)
    x_scaler_path = os.path.join(save_dir, 'x_scaler.pkl')
    y_scaler_path = os.path.join(save_dir, 'y_scaler.pkl')
    X, y, indices = create_sliding_window(data)
    x_scaler = joblib.load(x_scaler_path)
    y_scaler = joblib.load(y_scaler_path)
    X = x_scaler.transform(X.reshape(-1, X.shape[-1])).reshape(X.shape)
    y = y_scaler.transform(y.reshape(-1, 1))
    return (
        torch.tensor(X, dtype=torch.float32),
        torch.tensor(y, dtype=torch.float32),
        y_scaler
    )

def load_npy_from_folder(folder_path):
    all_data = []
    for file in os.listdir(folder_path):
        if file.endswith('.npy'):
            data = np.load(os.path.join(folder_path, file))
            all_data.append(data)
    return np.concatenate(all_data, axis=0)

def visualize_predictions(y_true, y_pred, save_path, mode):
    plt.figure(figsize=(24, 8))
    y_true = np.array(y_true).flatten()
    y_pred = np.array(y_pred).flatten()
    x = np.arange(len(y_true))
    plt.plot(x, y_true, color='blue', label='True Values', alpha=0.6)
    plt.plot(x, y_pred, color='red', label='Predicted Values', alpha=0.6)
    plt.xlabel('Sample')
    plt.ylabel('Value')
    plt.title(f'Prediction vs True ({mode})')
    plt.legend()
    plt.tight_layout()
    plt.savefig(os.path.join(save_path, f'{mode}_scatter_predictions-1.png'))
    plt.close()

def benchmark_model(model, repeat=100):
    model.eval()
    x = torch.randn(1, 128, 15).to(device)
    with torch.no_grad():
        import time
        start = time.time()
        for _ in range(repeat):
            _ = model(x)
        end = time.time()
    print(f"Avg inference time over {repeat} runs: {(end - start)/repeat*1000:.2f} ms")

def main():
    # 路径自行修改
    save_dir = "/home/scuee_user06/myh/轻量化/code/model_pruning-512"
    os.makedirs(save_dir, exist_ok=True)
    amount = 0.3  # 剪枝比例
    # filter剪枝保存的文件名/路径
    slim_model_path = os.path.join(save_dir, f"cnn_model_slim.pth")
    idx1 = np.load(os.path.join(save_dir, 'idx1.npy'))
    idx2 = np.load(os.path.join(save_dir, 'idx2.npy'))

    # 加载模型
    model = SlimModelCNN(len(idx1), len(idx2)).to(device)
    load_model(model, slim_model_path)
    model.eval()

    # 加载 scaler、测试数据
    y_scaler = joblib.load(os.path.join(save_dir, 'y_scaler.pkl'))
    test_data = load_npy_from_folder("/home/scuee_user06/myh/轻量化/飞行数据/processed_test-new")
    X_test, y_test, _ = preprocess_data(test_data, save_dir, is_train=False)

    # 推理与评估
    with torch.no_grad():
        X_test = X_test.to(device)
        y_pred = model(X_test).cpu().numpy()
        if y_pred.ndim == 3:
            y_pred = y_pred[:, -1, 0]
        elif y_pred.ndim == 2 and y_pred.shape[1] > 1:
            y_pred = y_pred[:, -1]
        else:
            y_pred = y_pred.flatten()
        y_true = y_test.cpu().numpy().squeeze()
        # 反归一化
        y_pred = y_scaler.inverse_transform(y_pred.reshape(-1, 1)).flatten()
        y_true = y_scaler.inverse_transform(y_true.reshape(-1, 1)).flatten()

        mse = mean_squared_error(y_true, y_pred)
        rmse = np.sqrt(mse)
        mae = mean_absolute_error(y_true, y_pred)
        r2 = r2_score(y_true, y_pred)
        print(f"[Eval Accuracy]")
        print(f"  MSE  = {mse:.6f}")
        print(f"  RMSE = {rmse:.6f}")
        print(f"  MAE  = {mae:.6f}")
        print(f"  R²   = {r2:.6f}")

        visualize_predictions(y_true, y_pred, save_dir, f'eval-剪枝{amount}-filter-1')
        print(f"Sample of Predicted Values: {y_pred[:7]}")
        print(f"Sample of True Values: {y_true[:7]}")

        benchmark_model(model)

        df = pd.DataFrame({'True': y_true, 'Pred': y_pred})
        df.to_csv(os.path.join(save_dir, 'predictions—1.csv'), index=False)

if __name__ == "__main__":
    main()
